import type { Plugin } from 'vite'
import * as parser from '@babel/parser'
import traverse from '@babel/traverse'
import generate from '@babel/generator'
import * as t from '@babel/types'

export default function removeTodoAst(): Plugin {
  return {
    name: 'vite-plugin-remove-todo-ast',
    enforce: 'pre',
    transform(code, id) {
      if (!id.match(/\.(js|ts|vue)$/)) return
      if (process.env.NODE_ENV !== 'development') return

      const ast = parser.parse(code, {
        sourceType: 'module',
        plugins: ['jsx', 'typescript']
      })

      traverse(ast, {
        enter(path) {
          // 移除 leadingComments 中含有 TODO: 和 TODO: end 的节点
          const comments = (path.node.leadingComments || [])
          let removeThis = false
          let hasTodoStart = false
          let hasTodoEnd = false

          comments.forEach(c => {
            if (/TODO:/i.test(c.value)) hasTodoStart = true
            if (/TODO:\s*end/i.test(c.value)) hasTodoEnd = true
          })

          if (hasTodoStart && hasTodoEnd) {
            removeThis = true
          }

          if (removeThis) {
            path.remove()
          }
        }
      })

      const output = generate(ast, {}, code)
      return { code: output.code, map: null }
    }
  }
}
